#!/usr/bin/env python
import sys
import os
import time
import platform
from racecar.parse import PosteriorExtractor
from racecar.bayes import HungryPosterior
from commands import getoutput
from optparse import OptionParser
from getpass import getpass

# Script for running Racecar in a parallel fashion.
# Supports many of the same command line arguments as racecar, and a few
# xgrid specific ones.

# Requires a Mac OS X system.
if platform.system() == "Darwin":
  import xg
else:
  print "Must be running on Darwin platform to execute this script.";
  sys.exit(1);

# Parse script arguments. 
def getSettings():
  # Create our option parser. 
  parser = OptionParser();
  
  # Add racecar options.
  parser.add_option("-i","--in",dest="datafile_racecar",
                    help="The location of the input Racecar format data file.");
  
  parser.add_option("","--orange",dest="datafile_orange",
                    help="The location of the input Orange format data file.");
  
  parser.add_option("--exhaustive",action="store_true",
                    dest="useExhaustive",default=True,
                    help="Use exhaustive searcher.");
  
  parser.add_option("--take_k",type="int",dest="takek",default=1,
                    help="Combination parameter for exhaustive search.");
  
  parser.add_option("-s","--postsize",dest="postSize",type="int",default=100,
                    help="Posterior size, per unit.");
  
  defaultProjectName = '-'.join([str(i) for i in time.localtime()[0:6]]);
  parser.add_option("--name",dest="projectName",default=defaultProjectName,
                    help="Project name - script will write output to files \
                    prefixed with the project name.");
  
  # Grid options.
  parser.add_option("--gid",dest="gid",type="int",default=0,
                    help="Grid ID. Defaults to 0.");
  
  parser.add_option("--hostname",dest="hostname",
                    help="Hostname of Xgrid.");
  
  parser.add_option("-u","--units",dest="units",type="int",
                    help="Number of grid units to use.");
  
  # Parse arguments.
  (options,args) = parser.parse_args();
  
  # Run sanity checks.
  if (options.datafile_orange == None) and (options.datafile_racecar == None):
    parser.error("Data file location is required, use -i or --orange.");
  
  if options.datafile_orange != None:
    checkFile(options.datafile_orange);
  else:
    checkFile(options.datafile_racecar);
    
  if options.hostname == None:
    parser.error("Xgrid hostname required, see usage using -h.");
  
  if options.units == None:
    parser.error("Number of grid units required, see usage using -h.");
  
  if options.takek < 0:
    parser.error("TakeK value must be greater than zero.");
  
  if options.units < 0:
    parser.error("Number of grid units must be positive and nonzero.");
  
  # Everything panned out, return the options/args.
  return (options,args);

# Check a file to make sure it exists, if not - throw an exception.
def checkFile(file):
  if not os.path.isfile(file):
    raise IOError, "\"" + str(file) + "\" cannot be found. Check your path.";
  
# Find the racecar binary location on the local system.
def findRacecar():
  # Racecar binary location.
  racecar = None;
  if platform.system() == "Windows":
    if os.path.isfile(r"..\bin\Racecar.exe"):
      racecar = r"..\bin\Racecar.exe";
    else:
      # Hope that it is on the path. 
      racecar = r"Racecar.exe";
  elif platform.system() == "Darwin":
    if os.path.isfile(r'../bin/racecar'):
      racecar = r"../bin/racecar";
    else:
      # Hope that it is on the path. 
      racecar = getoutput('which racecar');
  
  return racecar;

# Sets up the xgrid for use - returns a grid object.
# Requires an options object with hostname set. 
def setupGrid(options):
  # We need the xgrid password..
  print "Provide xgrid password..";
  password = getpass();
  
  # Do our xgrid jazz.
  print "Connecting to xgrid..";
  try:
    conn = xg.Connection(options.hostname,password);
    contr = xg.Controller(conn);
    grid = contr.grid(options.gid);
  except:
    print "Connection failed, this is probably due to an incorrect password."
    print "The actual (and probably very cryptic) error thrown was:"
    print "";
    raise;

  return grid;

# Returns an open file object with results from the job.
def getResults(options,job):
  # Check every 10 seconds to see if the job is done. 
  while not isDone(job):
    time.sleep(10);
    
  # Grab the job results, save it.
  filename = options.projectName + ".raw_results.tab";
  job.results(stdout=filename);
  
  # Return the filename.
  return filename;

# Returns true if the job is finished, false otherwise. 
def isDone(job):
  status = str(job.info()['jobStatus']);
  if status == 'Finished':
    return True;
  else:
    return False;
  
def createArgString(options,unit):
  argList = list();
  if options.useExhaustive:
    argList.append("--exhaustive");
    argList.append("--take_k " + str(options.takek));
    if options.datafile_orange != None:
      argList.append("--orange " + os.path.split(options.datafile_orange)[1]);
    else:
      argList.append("-i " + os.path.split(options.datafile_racecar)[1]);
    argList.append("--unit " + str(unit));
    argList.append("--units " + str(options.units));
    argList.append("-s " + str(options.postSize));
  
  return " ".join(argList);

def submitJobSpec(grid,spec):
  print "Submitting job to grid..";
  job = grid.batch(spec);
  
  return job;

def createJobSpec(options):
  # Create base spec. 
  spec = xg.JobSpecification();
  
  # Setup the files to add to the job specification.
  racecar = findRacecar();
  if options.datafile_orange != None:
    datapath = options.datafile_orange;
  else:
    datapath = options.datafile_racecar;
  binaryName = os.path.split(racecar)[1];
  binaryPath = os.path.realpath(racecar);
  dataName = os.path.split(datapath)[1];
  dataPath = os.path.realpath(datapath);
  spec.addFile(binaryPath,binaryName,1);
  spec.addFile(dataPath,dataName,0);
  
  if options.useExhaustive:
    # Create each task.
    for i in xrange(1,options.units+1):
      spec.addTask(binaryName,createArgString(options,i));
    
  return spec;

def waitForJob(options,job):
  print "Waiting for job to finish.."
  try:
    resultFile = getResults(options,job);
  except KeyboardInterrupt:
    # Delete jobs from the grid if someone or something interrupts while waiting.
    print "";
    print "** Keyboard interrupt detected, deleting job.. **";
    job.delete();
    print "";
    sys.exit(1);
  except:
    # Delete jobs from the grid if someone or something interrupts while waiting.
    print "";
    print "-------------------------------------------------------";
    print " Unknown exception detected, job will remain on grid.. ";
    print " Be sure to delete it in the future! ";
    print "-------------------------------------------------------";
    print "";
    raise;
  
  return resultFile;

def parseResults(options,file):
  # Create the extractor. 
  extractor = PosteriorExtractor(HungryPosterior());
  
  # Open a file object. 
  f = open(file);
  
  # Parse it.
  print "Parsing results..";
  extractor.parseFileObject(f);
  f.close();
  
  # Write the results to the file specified in options.
  outFileName = options.projectName + ".posterior.tab";
  out = open(outFileName,'w');
  extractor.getPosterior().write(out);
  out.close();

# Run exhaustive search method using the Xgrid.
def runExhaustive(options):
  # Connect to the grid.
  grid = setupGrid(options);
  
  # Create our job specification. 
  spec = createJobSpec(options);
  
  # Submit spec to the grid.
  job = submitJobSpec(grid,spec);
  
  # Wait for job to complete.
  resultsFileName = waitForJob(options,job);
  
  # Parse the results.
  parseResults(options,resultsFileName);
  
  # Kill the job on the grid.
  job.delete();

def main():
  # Get our command line arguments.
  (options,args) = getSettings();
  
  # If exhaustive, run it.
  if options.useExhaustive == True:
    runExhaustive(options);

if __name__ == "__main__":
  main();